"""
记忆管理器模块
实现不同类型记忆的专门管理器
"""

import json
import asyncio
from typing import Dict, Any, List, Optional
from datetime import datetime, timedelta
import logging

from .memory_architecture import MemoryItem, MemoryType, MemoryPriority

logger = logging.getLogger(__name__)

class BaseMemoryManager:
    """基础记忆管理器"""
    
    def __init__(self, architecture):
        self.architecture = architecture
        self.memory_type = None
    
    async def store(self, memory_item: MemoryItem) -> bool:
        """存储记忆项"""
        return await self.architecture.store_memory(memory_item)
    
    async def retrieve(self, query: str, user_id: str, **kwargs) -> List[MemoryItem]:
        """检索记忆项"""
        return await self.architecture.retrieve_memories(
            query, user_id, [self.memory_type], **kwargs
        )
    
    async def cleanup(self, user_id: str) -> int:
        """清理过期记忆"""
        return 0

class ShortTermMemoryManager(BaseMemoryManager):
    """短期记忆管理器"""
    
    def __init__(self, architecture):
        super().__init__(architecture)
        self.memory_type = MemoryType.SHORT_TERM
    
    async def store_conversation_turn(
        self, 
        user_message: str, 
        ai_response: str, 
        user_id: str, 
        session_id: str,
        metadata: Dict[str, Any] = None
    ) -> bool:
        """存储对话轮次"""
        try:
            # 存储用户消息
            user_memory = MemoryItem(
                id=f"user_{session_id}_{datetime.now().timestamp()}",
                content=user_message,
                memory_type=MemoryType.SHORT_TERM,
                priority=MemoryPriority.MEDIUM,
                user_id=user_id,
                session_id=session_id,
                timestamp=datetime.now(),
                metadata={
                    "role": "user",
                    "turn_type": "conversation",
                    **(metadata or {})
                }
            )
            
            # 存储AI回复
            ai_memory = MemoryItem(
                id=f"ai_{session_id}_{datetime.now().timestamp()}",
                content=ai_response,
                memory_type=MemoryType.SHORT_TERM,
                priority=MemoryPriority.MEDIUM,
                user_id=user_id,
                session_id=session_id,
                timestamp=datetime.now(),
                metadata={
                    "role": "assistant",
                    "turn_type": "conversation",
                    **(metadata or {})
                }
            )
            
            # 并发存储
            results = await asyncio.gather(
                self.store(user_memory),
                self.store(ai_memory),
                return_exceptions=True
            )
            
            return all(isinstance(r, bool) and r for r in results)
            
        except Exception as e:
            logger.error(f"Failed to store conversation turn: {e}")
            return False
    
    async def get_recent_context(
        self, 
        user_id: str, 
        session_id: str, 
        max_turns: int = 10
    ) -> List[Dict[str, str]]:
        """获取最近的对话上下文"""
        try:
            # 构建查询键
            key = self.architecture.get_memory_key(
                MemoryType.SHORT_TERM, user_id, session_id
            )
            
            # 从Redis获取所有记忆
            memory_data = await self.architecture.redis_client.hgetall(key)
            
            # 解析并排序
            memories = []
            for memory_id, memory_json in memory_data.items():
                try:
                    memory_dict = json.loads(memory_json)
                    memory_item = MemoryItem.from_dict(memory_dict)
                    memories.append(memory_item)
                except Exception as e:
                    logger.warning(f"Failed to parse memory {memory_id}: {e}")
            
            # 按时间排序
            memories.sort(key=lambda x: x.timestamp)
            
            # 转换为对话格式
            context = []
            for memory in memories[-max_turns*2:]:  # 每轮对话包含用户和AI消息
                role = memory.metadata.get("role", "user")
                context.append({
                    "role": role,
                    "content": memory.content,
                    "timestamp": memory.timestamp.isoformat()
                })
            
            return context[-max_turns*2:]  # 限制返回数量
            
        except Exception as e:
            logger.error(f"Failed to get recent context: {e}")
            return []
    
    async def cleanup(self, user_id: str) -> int:
        """清理过期的短期记忆"""
        try:
            cutoff_time = datetime.now() - timedelta(
                hours=self.architecture.config.SHORT_TERM_TTL_HOURS
            )
            
            # 获取所有短期记忆键
            pattern = f"memory:{MemoryType.SHORT_TERM.value}:{user_id}:*"
            keys = await self.architecture.redis_client.keys(pattern)
            
            cleaned_count = 0
            for key in keys:
                memory_data = await self.architecture.redis_client.hgetall(key)
                
                for memory_id, memory_json in memory_data.items():
                    try:
                        memory_dict = json.loads(memory_json)
                        memory_time = datetime.fromisoformat(memory_dict["timestamp"])
                        
                        if memory_time < cutoff_time:
                            await self.architecture.redis_client.hdel(key, memory_id)
                            cleaned_count += 1
                            
                    except Exception as e:
                        logger.warning(f"Failed to process memory {memory_id}: {e}")
            
            logger.info(f"Cleaned {cleaned_count} expired short-term memories for user {user_id}")
            return cleaned_count
            
        except Exception as e:
            logger.error(f"Failed to cleanup short-term memories: {e}")
            return 0

class LongTermMemoryManager(BaseMemoryManager):
    """长期记忆管理器"""
    
    def __init__(self, architecture):
        super().__init__(architecture)
        self.memory_type = MemoryType.LONG_TERM
    
    async def create_summary(
        self, 
        conversation_history: List[MemoryItem], 
        user_id: str
    ) -> Optional[MemoryItem]:
        """创建对话摘要"""
        try:
            if len(conversation_history) < 5:  # 太短的对话不需要摘要
                return None
            
            # 准备摘要内容
            conversation_text = self._format_conversation_for_summary(conversation_history)
            
            # 调用LLM生成摘要
            from langchain_core.prompts import ChatPromptTemplate
            from langchain_openai import ChatOpenAI
            import os
            
            model = ChatOpenAI(
                model=os.getenv("BASE_MODEL", "gpt-4"),
                api_key=os.getenv("OPENAI_API_KEY"),
                base_url=os.getenv("OPENAI_BASE_URL")
            )
            
            summary_prompt = ChatPromptTemplate.from_template("""
请对以下对话进行摘要，提取关键信息：

对话内容：
{conversation}

摘要要求：
1. 保留重要的事实信息
2. 记录用户的偏好和特征
3. 总结主要话题和结论
4. 保持简洁但信息完整

请用第一人称"我"的角度进行摘要，格式：
摘要内容 | 关键信息: 关键词1, 关键词2, 关键词3
""")
            
            response = await model.ainvoke(
                summary_prompt.format(conversation=conversation_text)
            )
            
            # 创建摘要记忆项
            summary_memory = MemoryItem(
                id=f"summary_{user_id}_{datetime.now().timestamp()}",
                content=response.content,
                memory_type=MemoryType.LONG_TERM,
                priority=MemoryPriority.HIGH,
                user_id=user_id,
                session_id="summary",
                timestamp=datetime.now(),
                metadata={
                    "type": "conversation_summary",
                    "original_count": len(conversation_history),
                    "time_range": {
                        "start": conversation_history[0].timestamp.isoformat(),
                        "end": conversation_history[-1].timestamp.isoformat()
                    }
                }
            )
            
            # 存储摘要
            success = await self.store(summary_memory)
            if success:
                self.architecture.metrics.summary_operations += 1
                logger.info(f"Created summary for {len(conversation_history)} messages")
                return summary_memory
            
            return None
            
        except Exception as e:
            logger.error(f"Failed to create summary: {e}")
            return None
    
    def _format_conversation_for_summary(self, conversation_history: List[MemoryItem]) -> str:
        """格式化对话用于摘要"""
        formatted = []
        for memory in conversation_history:
            role = memory.metadata.get("role", "user")
            timestamp = memory.timestamp.strftime("%H:%M")
            formatted.append(f"[{timestamp}] {role}: {memory.content}")
        
        return "\n".join(formatted)
    
    async def get_relevant_summaries(
        self, 
        query: str, 
        user_id: str, 
        top_k: int = 3
    ) -> List[MemoryItem]:
        """获取相关的历史摘要"""
        try:
            # 检索相关摘要
            summaries = await self.retrieve(query, user_id, top_k=top_k)
            
            # 过滤摘要类型的记忆
            summary_memories = [
                m for m in summaries 
                if m.metadata.get("type") == "conversation_summary"
            ]
            
            return summary_memories
            
        except Exception as e:
            logger.error(f"Failed to get relevant summaries: {e}")
            return []

# UserProfileManager 已移除 - 简化记忆系统，专注于核心对话记忆功能

class ContextManager(BaseMemoryManager):
    """上下文管理器"""
    
    def __init__(self, architecture):
        super().__init__(architecture)
        self.memory_type = MemoryType.CONTEXT
    
    async def store_context(
        self, 
        context_data: Dict[str, Any], 
        user_id: str, 
        session_id: str
    ) -> bool:
        """存储上下文信息"""
        try:
            context_memory = MemoryItem(
                id=f"context_{session_id}_{datetime.now().timestamp()}",
                content=json.dumps(context_data, ensure_ascii=False),
                memory_type=MemoryType.CONTEXT,
                priority=MemoryPriority.MEDIUM,
                user_id=user_id,
                session_id=session_id,
                timestamp=datetime.now(),
                metadata={
                    "type": "context_data",
                    "context_keys": list(context_data.keys())
                }
            )
            
            return await self.store(context_memory)
            
        except Exception as e:
            logger.error(f"Failed to store context: {e}")
            return False
    
    async def get_current_context(
        self, 
        user_id: str, 
        session_id: str
    ) -> Dict[str, Any]:
        """获取当前上下文"""
        try:
            # 获取最近的上下文记忆
            key = self.architecture.get_memory_key(
                MemoryType.CONTEXT, user_id, session_id
            )
            
            memory_data = await self.architecture.redis_client.hgetall(key)
            
            # 找到最新的上下文
            latest_context = {}
            latest_timestamp = None
            
            for memory_id, memory_json in memory_data.items():
                try:
                    memory_dict = json.loads(memory_json)
                    memory_timestamp = datetime.fromisoformat(memory_dict["timestamp"])
                    
                    if latest_timestamp is None or memory_timestamp > latest_timestamp:
                        latest_timestamp = memory_timestamp
                        latest_context = json.loads(memory_dict["content"])
                        
                except Exception as e:
                    logger.warning(f"Failed to parse context memory {memory_id}: {e}")
            
            return latest_context
            
        except Exception as e:
            logger.error(f"Failed to get current context: {e}")
            return {}
